import os
import base64
from datetime import datetime
from typing import List, Dict, Optional

import requests

GEN_PARAMS = {
    "temperature": 0.7,
    "modalities": ["text", "image"]
}
def _init_output_dir(dir_path: str) -> str:
    try:
        os.makedirs(dir_path, exist_ok=True)
    except Exception:
        pass
    return os.path.abspath(dir_path)


def _safe_stem(text: str, max_len: int = 40) -> str:
    s = (text or '').strip().replace('\n', ' ').replace('\r', ' ')
    for ch in ['\\', '/', ':', '*', '?', '"', '<', '>', '|', ',']:
        s = s.replace(ch, ' ')
    s = ' '.join(s.split())
    # Avoid spaces in filenames for better portability (pandoc/docx)
    s = s.replace(' ', '_')
    if len(s) > max_len:
        s = s[:max_len]
    if not s:
        s = 'img'
    return s


def download_image(image_url: str, save_dir: str, filename_hint: str = '') -> Optional[str]:
    """下载远程图片到本地，返回保存的绝对路径；失败返回 None。"""
    try:
        _init_output_dir(save_dir)
        headers = {"User-Agent": "Mozilla/5.0 (Downloader)"}
        resp = requests.get(image_url, headers=headers, timeout=60, stream=True)
        resp.raise_for_status()

        ctype = resp.headers.get('Content-Type', '').lower()
        ext = '.jpg'
        if 'png' in ctype:
            ext = '.png'
        elif 'webp' in ctype:
            ext = '.webp'
        elif 'jpeg' in ctype or 'jpg' in ctype:
            ext = '.jpg'

        ts = datetime.now().strftime('%Y%m%d%H%M%S')
        stem = _safe_stem(filename_hint) or 'img'
        save_name = f"doubao_{stem}_{ts}{ext}"
        save_path = os.path.abspath(os.path.join(save_dir, save_name))

        with open(save_path, 'wb') as f:
            for chunk in resp.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
        return save_path
    except Exception as e:
        print(f"下载图片失败：{e}")
        return None


def remove_watermark_with_banana_nano(input_image_path: str, output_image_path: str,
                                       model: Optional[str] = None) -> Optional[str]:
    """
    使用 ChatGPT（banana nano）进行去水印处理。

    - 优先使用环境变量 BANANA_NANO_MODEL 指定模型；否则默认 'banana-nano'，失败回退 'gpt-image-1'。
    - API Key 优先读取环境变量 AIHUBMIX_API_KEY；否则尝试从 Doubao 文件内的 API_KEY（若导入可用）。
    - 通过 OpenAI 兼容 images.edits 接口输出 b64 图片，落盘保存。

    返回：输出图片绝对路径；失败返回 None。
    """
    try:
        from openai import OpenAI  # 延迟导入，避免未使用时报错
    except Exception as e:
        print(f"缺少 openai 依赖或导入失败：{e}")
        return None

    # 获取 API Key：优先环境变量，其次复用 Doubao 的 API_KEY（若可导入）
    api_key = "sk-PjfrbJ5yZ1kAeSH6875c4fD5551f4f4b8e839d28Da91A6Db"
    if not api_key:
        print('未找到可用的 API Key（请设置环境变量 AIHUBMIX_API_KEY 或在 doubao 文件中配置 API_KEY）。')
        return None

    base_url = 'https://aihubmix.com/v1'
    client = OpenAI(api_key=api_key, base_url=base_url)

    # 将输入图片转换为base64格式
    try:
        with open(input_image_path, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
    except Exception as e:
        print(f"❌ 读取输入图片失败：{str(e)}")
        return None

    full_prompt = "将下面的图片去除ai生成的水印，保持质量。"
    try:
        print(f"\n📌 去水印...")
        # 调用API
        response = client.chat.completions.create(
            model="gemini-2.5-flash-image-preview",
            messages=[{"role": "user", "content": [
                {
                    "type": "text",
                    "text": full_prompt
                },
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}
                }
            ]}],
            **GEN_PARAMS
        )

        # 解析图片数据
        message = response.choices[0].message
        if not hasattr(message, "multi_mod_content") or not message.multi_mod_content:
            print(f"❌ 无multi_mod_content字段")
            return None

        # 查找有效Base64数据
        image_base64 = None
        for idx, mod_item in enumerate(message.multi_mod_content):
            if "inline_data" in mod_item and mod_item["inline_data"]:
                inline_data = mod_item["inline_data"]
                if "data" in inline_data and inline_data["data"].strip():
                    image_base64 = inline_data["data"].strip()
                    print(f"✅ 第{idx + 1}个元素找到图片数据（长度：{len(image_base64)}）")
                    break

        if not image_base64:
            print(f"❌ 去水印中未找到有效图片数据...")
            return None

        # 保存并返回绝对路径
        absolute_path = os.path.abspath(output_image_path)
        if base64_to_image(image_base64, absolute_path):
            print(f"✅ 生成成功：{absolute_path}")
            return absolute_path
    except Exception as e:
        print(f"❌ ai去水印生成异常：{str(e)}")
        return None

def base64_to_image(base64_str: str, save_path: str) -> bool:
    """Base64解码为图片并保存"""
    try:
        # 确保保存目录存在
        try:
            dir_path = os.path.dirname(save_path) or '.'
            os.makedirs(dir_path, exist_ok=True)
        except Exception:
            pass

        image_data = base64.b64decode(base64_str.strip())
        with open(save_path, "wb") as f:
            f.write(image_data)
        return True
    except Exception as e:
        print(f"❌ 图片保存失败：{str(e)}")
        return False

def process_doubao_image_url(image_url: str, prompt: str = '',
                              raw_dir: str = './agent/doubao_image_output/raw',
                              clean_dir: str = './agent/doubao_image_output/clean') -> Optional[Dict[str, str]]:
    """
    对单个 Doubao URL 执行：下载 -> 去水印。
    返回：{'prompt': 原始描述, 'raw_path': 本地原图, 'clean_path': 去水印后路径}；失败返回 None。
    """
    raw_path = download_image(image_url, save_dir=raw_dir, filename_hint=prompt)
    if not raw_path:
        print("下载 Doubao 图片失败，跳过去水印。")
        return None

    stem = os.path.splitext(os.path.basename(raw_path))[0]
    clean_path = os.path.abspath(os.path.join(clean_dir, f"{stem}_no_wm.png"))

    clean_out = remove_watermark_with_banana_nano(raw_path, clean_path)
    if not clean_out:
        return None

    return {
        'prompt': prompt,
        'raw_path': raw_path,
        'clean_path': clean_out,
    }

def batch_generate_and_clean_images(
    prompts: List[str],
    raw_dir: Optional[str] = None,
    clean_dir: Optional[str] = None,
    as_relative: bool = False,
) -> List[Dict[str, str]]:
    """
    综合流程：
    - 使用 Doubao 生成图片（得到 URL）
    - 下载图片到本地
    - 调用 ChatGPT banana nano 去水印

    参数：
    - raw_dir: 原始下载图片保存目录（默认 ./agent/doubao_image_output/raw）
    - clean_dir: 去水印后图片保存目录（默认 ./agent/doubao_image_output/clean）
    - as_relative: 如果为 True，返回的 image_path 将转换为相对路径 'images/<文件名>'
                   （仅当 clean_dir 指向文章同级的 images 目录时使用）

    返回列表：[{ 'prompt': str, 'image_path': 路径 }]
    当 as_relative=False 时，image_path 为去水印后本地绝对路径；
    当 as_relative=True 时，image_path 为 'images/<文件名>' 相对路径。
    """
    try:
        from agent.doubao_generate_iamge import generate_single_image  # 复用原始生成函数
    except Exception as e:
        print(f"导入 Doubao 生成函数失败：{e}")
        return []

    results: List[Dict[str, str]] = []
    # 目录默认值处理
    raw_dir = raw_dir or './agent/doubao_image_output/raw'
    clean_dir = clean_dir or './agent/doubao_image_output/clean'
    total = len(prompts)
    success = 0
    print(f"\n🚀 生成 + 去水印流程启动（共{total}张）")
    for p in prompts:
        url = generate_single_image(p)
        if not url:
            continue
        item = process_doubao_image_url(url, prompt=p, raw_dir=raw_dir, clean_dir=clean_dir)
        if not item or not item.get('clean_path'):
            continue
        cleaned_path = item['clean_path']
        # 统一路径分隔符
        cleaned_path = cleaned_path.replace('\\', '/')
        if as_relative:
            # 仅返回相对路径到 images 目录（假设 clean_dir 即为该 images 目录）
            rel = f"images/{os.path.basename(cleaned_path)}"
            results.append({'prompt': p, 'image_path': rel})
        else:
            results.append({'prompt': p, 'image_path': cleaned_path})
        success += 1

    print(f"\n完成：成功 {success} / {total}")
    return results


if __name__ == '__main__':
    # 简单自测入口（可按需注释/启用）
    DEMO_PROMPTS = [
        "A close-up of a delicate gold chain necklace lying on a soft, neutral fabric surface, showcasing its clean lines and subtle shine, with natural light illuminating its texture.",
    ]
    out = batch_generate_and_clean_images(DEMO_PROMPTS)
    for i, item in enumerate(out, 1):
        print(f"{i}. 描述: {item['prompt']}\n   去水印图片: {item['image_path']}")

